package cn.coisini.model.common.search;

import lombok.extern.log4j.Log4j2;

import java.util.ArrayList;
import java.util.List;

/**
 * @author xiangshaw
 * Description:
 */
@Log4j2
public class Trie {
    // Trie 树的根节点
    private final TrieNode root;

    public Trie() {
        // 初始化根节点
        root = new TrieNode();
        // 根节点的值设为空格字符，表示根节点不存储实际字符
        root.value = ' ';
    }

    /**
     * 插入字符串到 Trie 树中
     */
    public void insert(String word) {
        // 从根节点开始
        TrieNode currentNode = root;
        // 遍历单词的每个字符
        for (int i = 0; i < word.length(); i++) {
            // 获取当前字符
            char c = word.charAt(i);
            // 如果当前字符不在当前节点的子节点中，则创建新的节点
            if (!currentNode.children.containsKey(c)) {
                // 创建一个新节点并添加到子节点中
                currentNode.children.put(c, new TrieNode(c));
            }
            // 更新当前节点为下一个字符对应的节点（移动到子节点）
            currentNode = currentNode.children.get(c);
            // 将单词添加到该节点的单词列表中
            currentNode.wordsFromNode.add(word);
        }
        // 标记当前节点为单词的结尾
        currentNode.isEndOfWord = true;
    }

    /**
     * 根据前缀查询 Trie 树，返回以该前缀开头的所有单词
     */
    public List<String> startWith(String prefix) {
        // 用于存储匹配的单词
        List<String> matches = new ArrayList<>();
        // 从根节点开始
        TrieNode currentNode = root;
        // 遍历前缀的每个字符
        for (int i = 0; i < prefix.length(); i++) {
            // 获取当前字符
            char c = prefix.charAt(i);
            // 如果当前字符不在当前节点的子节点中，则返回空列表
            if (!currentNode.children.containsKey(c)) {
                // 返回空列表
                return matches;
            }
            // 更新当前节点为下一个字符对应的节点（移动到子节点）
            currentNode = currentNode.children.get(c);
        }
        // 添加从当前节点开始的所有单词到匹配列表中
        matches.addAll(currentNode.wordsFromNode);
        // 返回匹配列表
        return matches;
    }



    public static void main(String[] args) {
        Trie t = new Trie();
        t.insert("萧肖");
        t.insert("萧小");
        t.insert("xiao小");
        // 匹配萧
        List<String> ret = t.startWith("萧");
        // [萧肖, 萧小]
        log.info(ret);
    }
}
